import numpy as np
import tqdm, json

class System():
    def __init__(self, generators, loads, lines, max_load_threshold = 1) -> None:
        self.nodes = lines.copy()
        self.nodes.update(loads)
        self.nodes.update(generators)
        self.generators=generators
        self.loads=loads
        self.lines=lines
        self.threshold = max_load_threshold

    def _find_line_recursive(self, source_id, target_id, path_buffer):
        """
        使用DFS寻找从source到target的路径，优先选择负载较小的路径
        :param source: 源节点
        :param target: 目标节点
        :param visited: 已访问节点
        :return: 路径和路径的最小剩余容量
        """

        if source_id == target_id:
            return [[target_id], [np.inf], [False]]
        elif source_id in path_buffer:
            return path_buffer[source_id]
        elif self.nodes[source_id].is_open():
            return [[],[],[]]

        path_buffer[source_id] = [[],[],[]] # 防止环路

        best_path = []
        best_path_load_rate = []
        best_path_is_keys = []

        # 针对当前节点计算负载率
        max_load = max(self.nodes[target_id].get_load_list()) # 负载
        cap_value = self.nodes[source_id].get_capacity() # 当前容量
        cur_load_rate = self.nodes[source_id].get_load_rate(max_load) # 当前节点的负载率

        min_load_rate = np.inf
        mean_load_rate = np.inf
        over_load_num = np.inf
        total_paths = []
        for node_id in self.nodes[source_id].get_to_list():
            # 如果当前节点不是目标节点, 同时节点存在上一级联通节点, 则进一步遍历上一级节点
            [next_best_path, next_load_rate, is_keys] = self._find_line_recursive(node_id, target_id, path_buffer)

            if len(next_best_path)<=0: # 该子节点无法路由到目标负载
                continue
            # 更新当前节点
            path = [source_id]+next_best_path 
            load_rate = [cur_load_rate]+next_load_rate
            is_keys = [False]+is_keys
            
            total_paths.append([path,load_rate])
            # 当前线路的最大负载率, 对于线路的选择, 应当选择负载率最小的线路
            cur_max_load_rate = max(load_rate[:-1]) # 去除负载
            cur_mean_load_rate = np.mean(load_rate[:-1])
            cur_over_load_num = (np.array(load_rate[:-1])>self.threshold).sum()
            need_update = False
            # 根据线路的最大负载值作为判断依据, 找到负载最低的线路

            if cur_max_load_rate < min_load_rate:
                need_update = True
            elif abs(cur_max_load_rate - min_load_rate)<0.01: 
                if len(path) < len(best_path):
                    need_update = True
                elif cur_mean_load_rate < mean_load_rate:
                    need_update = True
                elif cur_over_load_num < over_load_num:
                    need_update = True
            if need_update:
                best_path = path
                best_path_load_rate = load_rate
                best_path_is_keys = is_keys
                min_load_rate = cur_max_load_rate
                mean_load_rate = cur_mean_load_rate
                over_load_num = cur_over_load_num
            
        # 当前节点是否存在多条线路, 是则说明这是关键性节点
        if len(total_paths)>1:
            best_path_is_keys[0] = True 
        
        # 添加路径信息缓存
        path_buffer[source_id] = [best_path, best_path_load_rate, best_path_is_keys]
        return [best_path, best_path_load_rate, best_path_is_keys]

    def find_min_load_line(self, load_id):
        
        best_path = []
        best_load_rates = []
        best_is_keys = []
        has_redundant = 0

        min_load_rate = np.inf
        mean_load_rate = np.inf
        over_load_num = np.inf
        for generator_id in self.generators.keys():
            if self.nodes[generator_id].get_capacity() > 0:  # 如果发电机还有剩余电力
                path, full_load_rates, is_keys  = self._find_line_recursive(generator_id, # 从发电到负载
                                                                          load_id, 
                                                                          {},
                                                                          ) 
                if len(path)<=0:
                    continue

                # 添加根节点
                path=["root"]+path
                full_load_rates = [0]+full_load_rates
                is_keys = [False]+is_keys

                has_redundant+=1 # 备线数量

                line_load_rates = full_load_rates[1:-1] # 去除根节点和负载部分的值
                cur_max_load_rate = max(line_load_rates) # 去除负载
                cur_mean_load_rate = np.mean(line_load_rates)
                cur_over_load_num = (np.array(line_load_rates)>self.threshold).sum()

                need_update = False
                if cur_max_load_rate<min_load_rate:
                    need_update = True
                elif abs(cur_max_load_rate - min_load_rate) < 0.01:
                    if len(path) < len(best_path):
                        need_update = True
                    elif cur_mean_load_rate < mean_load_rate:
                        need_update = True
                    elif cur_over_load_num < over_load_num:
                        need_update = True
                if need_update:
                    best_path = path
                    best_load_rates = full_load_rates
                    best_is_keys = is_keys
                    min_load_rate = cur_max_load_rate
                    mean_load_rate = cur_mean_load_rate
                    over_load_num = cur_over_load_num
        if has_redundant>1:
            best_is_keys[0] = True

        return best_path, best_load_rates, best_is_keys

    def optimize_line(self, overload_id, overload_path, overload_rates, overload_iskeys):

        need_replace_points = [] # 统计负载过高的节点
        need_release_loads = [] # 统计这些负载过高的节点需要释放掉多少负载
        load_node = self.loads[overload_id]
        for point_id, load_rate in zip(overload_path[:-1], overload_rates[:-1]):
            if load_rate < self.threshold:
                continue
            need_replace_points.append(point_id)
            # 负载节点的负载-当前节点的剩余容量, 即需要释放的容量
            need_release_loads.append(max(load_node["load"]) - self.nodes[point_id].get_capacity()) 
        
        # 先将超载的节点强行挂载, 然后通过优化让超载消失
        self.add_load(overload_id,overload_path,overload_iskeys, hard = True)
        # 逐个取出负载过高节点中的负载, 查看他们是否存在冗余线路, 并进行替换
        for point_id in need_replace_points: # 需要被替换的高负载节点
            done = False
            node_ = self.nodes[point_id]
            user_ids = list(node_["users"].keys()) # 看看有哪些负载在使用这个节点
            for user_id in user_ids: # 遍历这些负载
                best_path_org, is_keys_org = self.loads[user_id]["best_path"] # 查看这个负载目前使用的最优选路
                end_idx = best_path_org.index(point_id) # 找到这个高负载节点在这条最优线路中位置

                if not True in is_keys_org[:end_idx]: # 看看这个高负载节点前面有没有冗余线路, 有就代表可能有备用线路
                    # print(user_id," can't replace")
                    continue
                
                # 将原有负载释放, 然后重新寻找
                self.release_load(user_id, best_path_org)
                best_path_, loads_rate_, is_keys_ = self.find_min_load_line(user_id)
                if len(best_path_)>0 and max(loads_rate_[:-1])<self.threshold: # 找到了替代线路
                    self.add_load(user_id,best_path_,is_keys_)
                else: # 没有找到, 恢复原线路
                    self.add_load(user_id,best_path_org, is_keys_org, hard= True)
                # 检查节点是否还超载, 如果不超载了说明替换结束
                if node_.get_capacity()>0:
                    done = True
                    break
            
            if not done: # 当前节点无法被优化, 说明存在共线线路负载超标且共线线路均为唯一线路, 恢复线路
                self.release_load(overload_id, overload_path)
                return False
        return True

    def add_load(self, load_id, path, is_keys, hard = False):
        load_value = max(self.nodes[load_id].get_load_list())
        for node_id in path[1:-1]:
            node_ = self.nodes[node_id]
            node_.add_load(load_id, load_value, hard)
        self.loads[load_id].update_best_path(path,is_keys)
    
    def release_load(self, load_id, path):
        for node_id in path[1:-1]:
            node_ = self.nodes[node_id]
            node_.release_load(load_id)
        self.loads[load_id]["best_path"] = []

    def building_map(self):
        overload_loads = []
        # 遍历负载, 构建线路
        for load_id, load in tqdm.tqdm(self.loads.items()):
            # 构建最优线路
            best_path, best_load_rates, best_is_keys = self.find_min_load_line(load_id)
            if len(best_path)<=0:
                overload_loads.append(load_id)

            min_load_rate = max(best_load_rates[:-1])
            # 负载过高, 优化线路
            if min_load_rate>self.threshold:
                ret = self.optimize_line(
                        overload_id = load_id,
                        overload_path = best_path,
                        overload_rates = best_load_rates,
                        overload_iskeys = best_is_keys
                        )
                if not ret: #说明优化不成功
                    overload_loads.append(load_id)
            else:
                # 更新负载
                self.add_load(load_id,best_path,best_is_keys)

        total_json = {
            "nodes":[],
            "overloads": overload_loads
        }
        for node_id, node in self.nodes.items():
            node_json = {
                "rdfId": node.get_id(),
                "from": node.get_from_list(),
                "open": node["is_open"],
                "to": node.get_to_list(),
                "users":node["users"],
                "type": node["node_type"],
                "path": node["best_path"]
            }
            total_json["nodes"].append(node_json)

        with open("result.json","w") as f:
            json.dump(total_json,f, indent=4)

            


if __name__ == "__main__":

    from utils.cus_data_io import buid_test_data, load_data
    # generators, loads, lines = buid_test_data()
    generators, loads, lines = load_data()

    sys = System(generators,loads,lines,0.8)
    sys.building_map()